from hapi.model import Model
from paddle.fluid.layers import elementwise_add as add, elementwise_mul as mul
from paddle.fluid.dygraph import Layer, Sequential, Linear

__all__ = [
    "PreActResNet"
]

from hapi.vision.models.cifar.layers import Conv2d, BNAct, Pool2d, GlobalAvgPool, flatten


class SELayer(Layer):

    def __init__(self, in_channels, reduction):
        super().__init__()
        channels = in_channels // reduction
        self.pool = GlobalAvgPool()
        self.fc = Sequential(
            Linear(in_channels, channels, act='relu'),
            Linear(channels, in_channels, act='sigmoid'),
        )

    def forward(self, x):
        s = self.pool(x)
        s = flatten(s)
        s = self.fc(s)
        x = mul(x=x, y=s, axis=0)
        return x


class Shortcut(Sequential):

    def __init__(self, in_channels, out_channels, stride):
        layers = []
        if in_channels != out_channels or stride != 1:
            if stride != 1:
                layers.append(
                    Pool2d(2, 2, type='avg')
                )
            layers.append(Conv2d(in_channels, out_channels, 1, stride=1))
        super().__init__(*layers)


class BaiscBlock(Layer):
    def __init__(self, in_channels, out_channels, use_se):
        super().__init__()
        self.use_se = use_se
        self.bn_act1 = BNAct(in_channels, 'relu')
        self.conv1 = Conv2d(in_channels, out_channels, kernel_size=3)
        self.bn_act2 = BNAct(out_channels, 'relu')
        self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
        if use_se:
            self.se = SELayer(out_channels, reduction=8)

    def forward(self, x):
        identity = x
        x = self.bn_act1(x)
        x = self.conv1(x)
        x = self.bn_act2(x)
        x = self.conv2(x)
        if self.use_se:
            x = self.se(x)
        return add(identity, x)


class DownBlock(Layer):
    def __init__(self, in_channels, out_channels, stride=1, use_se=False):
        super().__init__()
        self.use_se = use_se
        self.bn_act1 = BNAct(in_channels, 'relu')
        self.conv1 = Conv2d(in_channels, out_channels, kernel_size=3, stride=stride)
        self.bn_act2 = BNAct(out_channels, 'relu')
        self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
        if self.use_se:
            self.se = SELayer(out_channels, reduction=8)

        self.shortcut = Shortcut(in_channels, out_channels, stride)

    def forward(self, x):
        x = self.bn_act1(x)
        identity = x
        x = self.conv1(x)
        x = self.bn_act2(x)
        x = self.conv2(x)
        if self.use_se:
            x = self.se(x)
        return add(self.shortcut(identity), x)


def _make_layer(in_channels, out_channels, num_units, stride, use_se):
    layers = [DownBlock(in_channels, out_channels, stride=stride, use_se=use_se)]
    for i in range(1, num_units):
        layers.append(BaiscBlock(out_channels, out_channels, use_se))
    return Sequential(*layers)


class PreActResNet(Model):

    def __init__(self, depth, k, use_se, num_classes=10):
        super().__init__()
        self.depth = depth
        self.k = k
        self.use_se = use_se
        self.num_classes = num_classes

        depth = self.depth

        ns = [(depth - 4) // 6] * 3

        stem_channels = 16
        channels_per_stage = [16, 32, 64]
        cs = [stem_channels] + [c * k for c in channels_per_stage]

        self.conv = Conv2d(3, 16, kernel_size=3)

        self.stage1 = _make_layer(cs[0], cs[1], ns[0], stride=1, use_se=use_se)
        self.stage2 = _make_layer(cs[1], cs[2], ns[1], stride=2, use_se=use_se)
        self.stage3 = _make_layer(cs[2], cs[3], ns[2], stride=2, use_se=use_se)

        out_channels = cs[-1]
        self.post_act = BNAct(out_channels, 'relu')
        self.pool = GlobalAvgPool()
        self.fc = Linear(out_channels, num_classes)

    def forward(self, x):
        x = self.conv(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)

        x = self.post_act(x)
        x = self.pool(x)
        x = flatten(x)
        x = self.fc(x)
        return x
